#!/usr/bin/env python3
"""
简化但改进的训练脚本
基于quick_train但加入train_helmet_detection的优势
"""

import os
import sys

# 修复OpenMP冲突
os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'
os.environ['OMP_NUM_THREADS'] = '1'

from ultralytics import YOLO
import torch
import warnings

warnings.filterwarnings('ignore', category=UserWarning)

def main():
    print("="*60)
    print("🛠️ 改进版安全帽检测训练")
    print("="*60)
    
    # GPU检测
    if torch.cuda.is_available():
        gpu_props = torch.cuda.get_device_properties(0)
        gpu_memory = gpu_props.total_memory / (1024**3)
        gpu_name = gpu_props.name
        
        print(f"✅ 检测到GPU: {gpu_name}")
        print(f"显存: {gpu_memory:.1f}GB")
        
        # 根据显存调整参数
        if gpu_memory <= 4:  # GTX 1050 Ti
            batch_size = 12
            workers = 2
            print("🎯 4GB显存优化配置")
        else:
            batch_size = 16
            workers = 4
            print("🎯 大显存配置")
            
        device = 'cuda'
    else:
        print("❌ 未检测到GPU，使用CPU")
        device = 'cpu'
        batch_size = 4
        workers = 2
    
    print(f"设备: {device}")
    print(f"批次大小: {batch_size}")
    print(f"工作线程: {workers}")
    
    # 检查数据集
    if not os.path.exists('data.yaml'):
        print("\n❌ 找不到data.yaml文件")
        return
    
    # 加载模型 - 尝试YOLOv8s获得更好效果
    print("\n📦 加载YOLOv8s模型...")
    try:
        model = YOLO('yolov8s.pt')  # 使用更大的模型
        print("✅ YOLOv8s模型加载成功")
    except Exception as e:
        print(f"❌ YOLOv8s加载失败，回退到YOLOv8n: {e}")
        model = YOLO('yolov8n.pt')
    
    print("\n🎯 开始训练...")
    
    try:
        # 训练参数 - 比quick_train更好，但比train_helmet_detection简单
        results = model.train(
            data='data.yaml',
            epochs=80,  # 比quick_train的50更多
            imgsz=640,
            batch=batch_size,
            device=device,
            workers=workers,
            project='runs/detect',
            name='helmet_s_improved',
            save=True,
            save_period=10,
            patience=20,  # 早停耐心值
            cache=False,
            augment=True,
            verbose=True,
            # 启用混合精度训练(如果是GPU)
            amp=True if device == 'cuda' else False,
        )
        
        print("\n🎉 训练完成!")
        print(f"最佳模型: runs/detect/helmet_s_improved/weights/best.pt")
        
        # 验证模型
        print("\n📊 验证模型...")
        metrics = model.val(device=device)
        print(f"mAP50: {metrics.box.map50:.3f}")
        print(f"mAP50-95: {metrics.box.map:.3f}")
        
        # 与之前的76.8%对比
        if metrics.box.map50 > 0.768:
            improvement = metrics.box.map50 - 0.768
            print(f"🚀 相比之前提升: +{improvement:.3f} ({improvement*100:.1f}%)")
        
    except Exception as e:
        print(f"\n❌ 训练失败: {e}")
        
        # 如果是显存问题，尝试降级
        if "memory" in str(e).lower() or "cuda" in str(e).lower():
            print("\n🔄 显存不足，尝试YOLOv8n + 小batch...")
            try:
                model_n = YOLO('yolov8n.pt')
                results = model_n.train(
                    data='data.yaml',
                    epochs=80,
                    imgsz=640,
                    batch=8,  # 更小的batch
                    device=device,
                    workers=2,
                    project='runs/detect',
                    name='helmet_n_safe',
                    save=True,
                    save_period=10,
                    patience=20,
                    cache=False,
                    augment=True,
                    verbose=True,
                    amp=True if device == 'cuda' else False,
                )
                print("✅ 降级训练成功!")
                
                metrics = model_n.val(device=device)
                print(f"mAP50: {metrics.box.map50:.3f}")
                print(f"mAP50-95: {metrics.box.map:.3f}")
                
            except Exception as e2:
                print(f"❌ 降级训练也失败: {e2}")

if __name__ == "__main__":
    main() 